# rnn 序列

from keras.backend import maximum, shape
from keras.engine.base_layer import Layer
from keras.models import Sequential
from keras import layers
import numpy as np
from numpy.core.defchararray import mod
import six.moves

class CharTable(object):
    def __init__(self, char) -> None:
        super().__init__()
        self.char = sorted(set(char))
        self.char_indices = dict((ch,i) for i,ch in enumerate(self.char))
        self.indices_char = dict((i,ch) for i,ch in enumerate(self.char))
    
    def encode(self, C, num_rows):
        x = np.zeros((num_rows,len(self.char)))
        for i,ch in enumerate(C):
            x[i, self.char_indices[ch]] = 1
        return x

    def decode(self,x, calc_argmax=True):
        if calc_argmax:
            x = x.argmax(axis=-1)
        return "".join(self.indices_char[x] for x in x)

class colors:
    ok = '\033[92m'
    fail = '\033[91m'
    cloase = '\033[0m'

TRAINING_SIZE = 50000
DIGITS = 3
REVERSE = True

MAXLEN = DIGITS + 1 + DIGITS
chars = '0123456789- '
ctable = CharTable(chars)

questions = []
expected = []
seen = set()
print('Generating data.....')
while len(questions) < TRAINING_SIZE:
    f = lambda: int(''.join(np.random.choice(list('0123456789'))
            for i in range(np.random.randint(1,DIGITS+1))))
    a, b = f(),f()
    key = tuple(sorted((a,b)))
    if key in seen:
        continue
    seen.add(key)
    q = '{}-{}'.format(a,b)
    query = q + ' ' * (MAXLEN - len(q))
    ans = str(a - b)
    ans += " "*(DIGITS + 1 -len(ans))
    if REVERSE:
        query = query[::-1]
    questions.append(query)
    expected.append(ans)
print("Total subtraction questions:",len(questions))

print('Vectorization:')
x = np.zeros((len(questions),MAXLEN,len(chars)),dtype=np.bool)
y = np.zeros((len(questions),DIGITS + 1,len(chars)),dtype=np.bool)
for i,sentence in enumerate(questions):
    x[i] = ctable.encode(sentence,MAXLEN)
for i,sentence in enumerate(expected):
    y[i] = ctable.encode(sentence,DIGITS + 1)

indices = np.arange(len(y))
np.random.shuffle(indices)
x = x[indices]
y = y[indices]

split_at = len(x) - len(x)//10
(x_train,x_val) = x[:split_at],x[split_at:]
(y_train,y_val) = y[:split_at],y[split_at:]
print("Training data shape:",x_train.shape)
print("val data shape:",x_val.shape)
print(x_train[:5])
print(y_train[:5])

RNN = layers.LSTM
HIDDEN_SIZE = 128
BATCH_SIZE = 128
LAYERS = 1

print("Build model:")
model = Sequential()
model.add(RNN(HIDDEN_SIZE,input_shape=(MAXLEN,len(chars))))
model.add(layers.RepeatVector(DIGITS + 1))
for _ in range(LAYERS):
    model.add(RNN(HIDDEN_SIZE,return_sequences=True))
model.add(layers.TimeDistributed(layers.Dense(len(chars))))
model.add(layers.Activation('softmax'))
model.compile(loss='categorical_crossentropy', optimizer='rmsprop', metrics=['accuracy'])

for iteration in range(1,200):
    print('-'*50)
    model.fit(x_train,y_train,batch_size=BATCH_SIZE,
            epochs=1,validation_data=(x_val,y_val))
    for i in range(10):
        ind = np.random.randint(0,len(x_val))
        rowx,rowy = x_val[np.array([ind])],y_val[np.array([ind])]
    preds = model.predict_classes(rowx, verbose=0)
    q = ctable.decode(rowx[0])
    correct = ctable.decode(rowy[0])
    guess = ctable.decode(preds[0],calc_argmax=False)
    print('Q',q[::-1] if REVERSE else q, end=' ')
    print('T',correct,end='')
    if correct == guess:
        print(colors.ok + 'yes' + colors.cloase,end=" ")
    else:
        print(colors.fail + 'no' + colors.cloase,end=" ")
        print(guess)